import pymysql as sql
import schedule
import time
from datetime import datetime
from api.zabbix_api import get_hosts_simple
from feishu.feishu_webhook import FeiShuWebhook
import hashlib
import json


class ZabbixAlertMonitor:
    def __init__(self):
        """初始化监控器"""
        # 数据库连接配置
        self.db_config = {
            'host': '10.104.14.27',
            'user': 'zabbix', 
            'password': 'Chery@zabbix_pwd',
            'database': 'zabbix'
        }
        
        # 飞书webhook配置
        self.webhook = FeiShuWebhook("https://open.feishu.cn/open-apis/bot/v2/hook/81852277-30cb-47a9-92be-b39543345b0e")
        
        # 用于存储已通知的告警，避免重复通知
        self.notified_alerts = set()
        
        # 连接数据库
        self.db_connection = None
        self.init_db_connection()
    
    def init_db_connection(self):
        """初始化数据库连接"""
        try:
            self.db_connection = sql.connect(**self.db_config)
            print("数据库连接成功")
        except Exception as e:
            print(f"数据库连接失败: {e}")
    
    def get_current_alerts(self):
        """获取当前所有告警"""
        try:
            self.init_db_connection()
            cursor = self.db_connection.cursor()
            query = """
            SELECT 
                h.hostid, 
                h.host, 
                h.name as hostname, 
                COUNT(DISTINCT p.eventid) as problem_count, 
                COUNT(DISTINCT a.alertid) as alert_count, 
                MAX(p.severity) as max_severity, 
                GROUP_CONCAT(DISTINCT p.name ORDER BY p.severity DESC SEPARATOR '; ') as problems,
                GROUP_CONCAT(DISTINCT concat(ht.tag, ":", ht.value) SEPARATOR ', ') as tags,
                GROUP_CONCAT(DISTINCT hg.groupid SEPARATOR ', ') as group_ids,
                GROUP_CONCAT(DISTINCT g.name SEPARATOR ', ') as host_groups
            FROM problem p
            INNER JOIN triggers t ON p.objectid = t.triggerid
            INNER JOIN functions f ON t.triggerid = f.triggerid
            INNER JOIN items i ON f.itemid = i.itemid
            INNER JOIN hosts h ON i.hostid = h.hostid
            LEFT JOIN host_tag ht ON i.hostid = ht.hostid
            LEFT JOIN hosts_groups hg ON h.hostid = hg.hostid
            LEFT JOIN hstgrp g ON hg.groupid = g.groupid
            LEFT JOIN events e ON p.eventid = e.eventid
            LEFT JOIN alerts a ON e.eventid = a.eventid
            WHERE p.r_eventid IS NULL 
                AND p.source = 0 
                AND p.object = 0
                AND g.name in ("Applications", "AP办公网段", "AP生产网段", "生产交换机网段", "办公交换机网段", "重要业务网段", "监控交换机网段")
            GROUP BY h.hostid, h.host, h.name
            ORDER BY problem_count DESC, max_severity DESC;
            """
            
            cursor.execute(query)
            results = cursor.fetchall()
            return results
            
        except Exception as e:
            print(f"获取告警数据失败: {e}")
            return []
    
    def get_host_tags_info(self, host_ip):
        """获取主机的标签信息"""
        try:
            host_info = get_hosts_simple(host_ip)
            if not host_info or len(host_info) == 0:
                return None
            
            host_data = host_info[0]
            tags = host_data.get('tags', [])
            
            # 提取设备名称和类型信息
            device_name = host_data.get('name', host_ip)  # 默认使用hostname，如果没有则使用IP
            type_tag = None
            position_tag = None
            
            for tag in tags:
                if tag.get('tag') == '类型':
                    type_tag = tag.get('value')
                elif tag.get('tag') == '位置':
                    position_tag = tag.get('value')
            
            return {
                'device_name': device_name,
                'type': type_tag,
                'position': position_tag,
                'tags': tags
            }
            
        except Exception as e:
            print(f"获取主机 {host_ip} 标签信息失败: {e}")
            return None
    
    def classify_alerts_by_type(self, alerts, new_alert_hashes=None):
        """按类型分类告警，包括新增告警分类"""
        classified_alerts = {}
        
        # 新增告警分类键
        NEW_ALERTS_KEY = "🚨 新增告警"
        
        for alert in alerts:
            host_id, host_ip, hostname, problem_count, alert_count, max_severity, problems, tags, group_ids, host_groups = alert
            
            # 生成当前告警的hash值
            alert_hash = self.generate_alert_hash(host_ip, problems)
            
            # 获取主机详细信息
            host_info = self.get_host_tags_info(host_ip)
            if not host_info:
                continue
            
            device_type = host_info['type'] or '未分类设备'
            device_name = host_info['device_name']
            position = host_info['position']
            
            # 构建主机信息（包括设备名称、位置和主机组信息）
            host_display_info = {
                'ip': host_ip,
                'hostname': hostname,
                'device_name': device_name,
                'position': position,
                'problem_count': problem_count,
                'problems': problems,
                'max_severity': max_severity,
                'tags': tags,
                'group_ids': group_ids,
                'host_groups': host_groups,
                'hash': alert_hash  # 存储hash值，便于后续处理
            }
            
            # 判断是否为新增告警
            is_new_alert = new_alert_hashes and alert_hash in new_alert_hashes
            
            # 如果是新增告警，只添加到新增告警分类
            if is_new_alert:
                if NEW_ALERTS_KEY not in classified_alerts:
                    classified_alerts[NEW_ALERTS_KEY] = []
                # 添加标记为新增的告警信息
                new_alert_info = host_display_info.copy()
                new_alert_info['is_new'] = True
                classified_alerts[NEW_ALERTS_KEY].append(new_alert_info)
            else:
                # 非新增告警才添加到普通分类
                if device_type not in classified_alerts:
                    classified_alerts[device_type] = []
                classified_alerts[device_type].append(host_display_info)
        
        return classified_alerts
    
    def generate_alert_hash(self, host_ip, problems):
        """生成告警唯一标识，用于去重和检测恢复"""
        content = f"{host_ip}:{problems}"
        return hashlib.md5(content.encode('utf-8')).hexdigest()
    
    def check_new_alerts(self):
        """检查是否有新告警和恢复的告警"""
        try:
            alerts = self.get_current_alerts()
            if not alerts:
                print(f"[{datetime.now()}] 当前没有活跃告警")
                # 如果之前有告警，现在没有了，说明所有告警都恢复了
                if self.notified_alerts:
                    print(f"[{datetime.now()}] 所有告警已恢复")
                    # 构建恢复通知
                    recovered_alerts = list(self.notified_alerts)
                    self.process_and_notify_alerts([], recovered_alerts)
                    # 清空已通知的告警集合
                    self.notified_alerts = set()
                return
            
            current_alert_hashes = set()
            alert_hash_map = {}  # 用于存储hash到告警信息的映射
            
            for alert in alerts:
                host_id, host_ip, hostname, problem_count, alert_count, max_severity, problems, tags, group_ids, host_groups = alert
                alert_hash = self.generate_alert_hash(host_ip, problems)
                current_alert_hashes.add(alert_hash)
                alert_hash_map[alert_hash] = alert
            
            # 检查是否有新的告警
            new_alerts = current_alert_hashes - self.notified_alerts
            
            # 检查是否有恢复的告警
            recovered_alerts = self.notified_alerts - current_alert_hashes
            
            # 如果是首次运行（notified_alerts为空），则所有当前告警都视为新告警
            is_first_run = len(self.notified_alerts) == 0
            
            if new_alerts or recovered_alerts or is_first_run:
                # 首次运行时，所有告警都标记为新告警
                if is_first_run:
                    print(f"[{datetime.now()}] 首次运行，发现 {len(current_alert_hashes)} 个告警，全部标记为新告警")
                else:
                    print(f"[{datetime.now()}] 发现 {len(new_alerts)} 个新告警，{len(recovered_alerts)} 个告警恢复")
                
                self.process_and_notify_alerts(alerts, recovered_alerts)
                
                # 更新已通知的告警集合
                self.notified_alerts = current_alert_hashes
            else:
                print(f"[{datetime.now()}] 没有新告警或恢复的告警，跳过通知")
                
        except Exception as e:
            print(f"检查告警时出错: {e}")
    
    def format_feishu_message(self, classified_alerts, recovered_alerts=None):
        """格式化飞书消息，包括新增告警和恢复告警"""
        # 过滤掉未分类设备
        filtered_alerts = {k: v for k, v in classified_alerts.items() if k != '未分类设备'}
        
        # 构建消息头部
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        message = f"🚨 **Zabbix告警监控报告**\n\n"
        message += f"📅 **检查时间**: {current_time}\n"
        
        # 统计总主机数和问题总数
        total_hosts = sum(len(hosts) for hosts in filtered_alerts.values())
        total_problems = sum(sum(host['problem_count'] for host in hosts) for hosts in filtered_alerts.values())
        message += f"📊 **告警类别**: {len(filtered_alerts)} 类\n"
        message += f"🏠 **受影响主机**: {total_hosts} 台\n"
        message += f"🔍 **问题总数**: {total_problems} 个\n\n"
        
        # 如果没有告警，直接返回消息
        if not filtered_alerts:
            message += "🎉 当前没有活跃告警，系统运行正常！\n\n"
            message += "_消息由 Zabbix 监控系统自动发送_"
            return message
        
        # 按类型分组显示告警 - 简化格式
        message += "## � 告警详情\n\n"
        
        # 定义自定义排序顺序
        custom_order = ['重要业务网段', '生产交换机网段']
        
        # 分离新增告警，单独处理
        new_alerts_key = "🚨 新增告警"
        new_alerts = filtered_alerts.pop(new_alerts_key, [])
        
        # 按照自定义顺序处理告警分类
        for category in custom_order:
            if category in filtered_alerts:
                hosts = filtered_alerts[category]
                message += f"� {category}\n"
                
                # 处理该类别的所有主机
                message += self._format_host_alerts(hosts) + "\n"
        
        # 处理剩余的其他分类
        for device_type, hosts in filtered_alerts.items():
            # 判断是否为生产网相关设备类型
            production_keywords = ['生产', 'prod', '生产网', '核心', '业务']
            is_production = any(keyword.lower() in device_type.lower() for keyword in production_keywords)
            
            # 根据类型选择圆点颜色
            if is_production:
                message += f"🔴 {device_type}\n"
            else:
                message += f"🟡 {device_type}\n"
            
            # 处理该类别的所有主机
            message += self._format_host_alerts(hosts)
            message += "\n"
        
        # 在最后显示新增告警信息
        if new_alerts:
            message += "## 🆕 新增告警\n\n"
            message += "🔴 🚨 新增告警\n"
            message += self._format_host_alerts(new_alerts)
            message += "\n"
        
        # 添加总结
        message += "## 📈 总结\n\n"
        message += f"🔍 本次检查发现 {len(filtered_alerts)} 种设备类型共 {total_hosts} 台主机存在 {total_problems} 个告警\n"
        
        # 在最后显示恢复告警信息
        if recovered_alerts and len(recovered_alerts) > 0:
            message += f"✅ 同时有 {len(recovered_alerts)} 个告警已恢复\n"
        
        message += "💡 请及时处理相关问题，确保系统稳定运行\n\n"
        message += "_消息由 Zabbix 监控系统自动发送_"
        
        return message
    
    def _format_host_alerts(self, hosts):
        """格式化主机告警信息，返回格式化的字符串"""
        formatted_lines = []
        for host in hosts:
            # 处理主机标签
            host_tags = []
            if host['position']:
                host_tags.append(f"位置:{host['position']}")
            if host['device_name'] != host['ip']:
                host_tags.append(f"设备:{host['device_name']}")
            
            # 如果没有标签，显示"无"
            tag_display = "; ".join(host_tags) if host_tags else "无"
            
            # 处理问题列表
            if host['problems']:
                problems_list = [p.strip() for p in host['problems'].split('; ') if p.strip()]
            else:
                problems_list = ["未知问题"]
            
            # 显示主机信息
            host_display = f"{host['ip']}"
            if host['device_name'] != host['ip']:
                host_display += f" ({host['device_name']})"
            
            # 显示每个问题
            for problem in problems_list:
                problem_display = f"{problem}"
                
                # 如果是新增告警，添加特殊标记
                if host.get('is_new'):
                    problem_display = "🆕 " + problem_display
                
                # 用逗号分隔三列
                formatted_lines.append(f"{problem_display}, {host_display}, {tag_display}")
        
        return "\n".join(formatted_lines)
    
    
    def get_severity_emoji(self, severity):
        """根据严重程度返回对应的表情符号"""
        severity_map = {
            0: "✅",  # Not classified
            1: "🟢",  # Information  
            2: "🟡",  # Warning
            3: "🟠",  # Average
            4: "🔴",  # High
            5: "🟣"   # Disaster
        }
        return severity_map.get(severity, "⚪")
    
    def process_and_notify_alerts(self, alerts, recovered_alerts=None):
        """处理告警并发送通知，包括恢复告警"""
        try:
            # 获取当前告警的hash集合
            current_alert_hashes = set()
            for alert in alerts:
                host_id, host_ip, hostname, problem_count, alert_count, max_severity, problems, tags, group_ids, host_groups = alert
                alert_hash = self.generate_alert_hash(host_ip, problems)
                current_alert_hashes.add(alert_hash)
            
            # 找出新增告警的hash集合
            # 如果是首次运行（notified_alerts为空），则所有当前告警都视为新告警
            is_first_run = len(self.notified_alerts) == 0
            if is_first_run:
                new_alert_hashes = current_alert_hashes
            else:
                new_alert_hashes = current_alert_hashes - (self.notified_alerts - (recovered_alerts or set()))
            
            # 分类告警
            classified_alerts = self.classify_alerts_by_type(alerts, new_alert_hashes)
            
            # 格式化消息
            message = self.format_feishu_message(classified_alerts, recovered_alerts)
            
            # 发送飞书通知
            print("正在发送飞书通知...")
            self.webhook.send(message)
            print("飞书通知发送完成")
            
        except Exception as e:
            print(f"处理告警时出错: {e}")
    
    def start_monitoring(self):
        """启动定时监控"""
        print("🚀 Zabbix告警监控服务启动")
        print(f"⏰ 监控间隔: 每分钟检查一次")
        print(f"🕐 启动时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("-" * 50)
        
        # 立即执行一次检查
        print("执行初始检查...")
        self.check_new_alerts()
        
        # 设置定时任务 - 每5秒执行一次（用于测试）
        schedule.every(5).seconds.do(self.check_new_alerts)
        
        # 保持程序运行
        try:
            while True:
                schedule.run_pending()
                time.sleep(5)  # 每5秒检查一次是否有待执行的任务
        except KeyboardInterrupt:
            print("\n🛑 监控服务已停止")
        except Exception as e:
            print(f"监控服务运行出错: {e}")
        finally:
            if self.db_connection:
                self.db_connection.close()
                print("数据库连接已关闭")


def main():
    """主函数"""
    monitor = ZabbixAlertMonitor()
    monitor.start_monitoring()


if __name__ == "__main__":
    main()